import numpy as np
from scipy.sparse import diags, csc_matrix, csr_matrix, lil_matrix, eye
from scipy.sparse.linalg import spsolve, bicgstab, cg

class MultiGridSolver:
    """多重网格求解器，用于解决线性方程组 Ax = b
    
    Attributes:
        max_levels: 最大网格层数
        pre_smoothing: 前向光滑次数
        post_smoothing: 后向光滑次数
        tolerance: 收敛容差
        max_iterations: 最大迭代次数
        min_iterations: 最小迭代次数，防止过早收敛
        residual_norm_type: 残差计算方式，'absolute'或'relative'
    """
    
    def __init__(self, max_levels=4, pre_smoothing=2, post_smoothing=2, tolerance=1e-6, max_iterations=100, min_iterations=1, residual_norm_type='absolute'):
        """初始化多重网格求解器
        
        Args:
            max_levels: 最大网格层数
            pre_smoothing: 前向光滑次数
            post_smoothing: 后向光滑次数
            tolerance: 收敛容差
            max_iterations: 最大迭代次数
            min_iterations: 最小迭代次数
            residual_norm_type: 残差计算方式
        """
        self.max_levels = max_levels
        self.pre_smoothing = pre_smoothing
        self.post_smoothing = post_smoothing
        self.tolerance = tolerance
        self.max_iterations = max_iterations
        self.min_iterations = min_iterations
        self.residual_norm_type = residual_norm_type
    
    def _create_restriction_operator(self, n_fine):
        """创建限制算子（使用稀疏矩阵）
        
        Args:
            n_fine: 细网格大小
            
        Returns:
            scipy.sparse.csr_matrix: 限制算子
        """
        n_coarse = (n_fine + 1) // 2
        R = lil_matrix((n_coarse, n_fine))
        
        # 使用稀疏矩阵构建限制算子
        for i in range(n_coarse):
            j = 2 * i
            # 边界点处理
            if j == 0:
                R[i, 0] = 0.75
                if j+1 < n_fine:
                    R[i, 1] = 0.25
            elif j >= n_fine - 1:
                if j-1 >= 0:
                    R[i, j-1] = 0.25
                if j < n_fine:
                    R[i, j] = 0.75
            else:
                # 内部点使用全加权限制
                R[i, j-1] = 0.25
                R[i, j] = 0.5
                if j+1 < n_fine:
                    R[i, j+1] = 0.25
        
        return R.tocsr()
    
    def _create_prolongation_operator(self, n_coarse, n_fine):
        """创建插值算子（使用稀疏矩阵）
        
        Args:
            n_coarse: 粗网格大小
            n_fine: 细网格大小
            
        Returns:
            scipy.sparse.csr_matrix: 插值算子
        """
        P = lil_matrix((n_fine, n_coarse))
        
        # 使用稀疏矩阵构建插值算子
        for i in range(n_fine):
            # 偶数点直接对应到粗网格
            if i % 2 == 0 and i//2 < n_coarse:
                P[i, i//2] = 1.0
            # 奇数点使用线性插值
            else:
                coarse_idx_left = (i-1)//2 if i > 0 else 0
                coarse_idx_right = coarse_idx_left + 1 if coarse_idx_left + 1 < n_coarse else coarse_idx_left
                
                if coarse_idx_left == coarse_idx_right:
                    P[i, coarse_idx_left] = 1.0
                else:
                    P[i, coarse_idx_left] = 0.5
                    P[i, coarse_idx_right] = 0.5
        
        return P.tocsr()
    
    def _jacobi_smoothing(self, A, x, b, num_iterations, omega=0.67):
        """雅可比光滑（适用于稀疏矩阵）
        
        Args:
            A: 稀疏系数矩阵
            x: 当前解
            b: 右端项
            num_iterations: 光滑次数
            omega: 松弛因子
            
        Returns:
            numpy.ndarray: 光滑后的解
        """
        # 提取对角线元素
        diag = A.diagonal()
        D_inv = 1.0 / diag
        
        # 对角线元素为零的处理
        D_inv[np.abs(diag) < 1e-10] = 0.0
        
        # 光滑迭代
        for _ in range(num_iterations):
            residual = b - A.dot(x)
            x += omega * D_inv * residual
        
        return x
    
    def _block_gauss_seidel(self, A, x, b, block_size=10, num_iterations=1):
        """分块高斯-赛德尔光滑（适用于带状矩阵）
        
        Args:
            A: 稀疏系数矩阵
            x: 当前解
            b: 右端项
            block_size: 分块大小
            num_iterations: 光滑次数
            
        Returns:
            numpy.ndarray: 光滑后的解
        """
        n = A.shape[0]
        A_dense = A.toarray() if n <= 1000 else None
        
        for _ in range(num_iterations):
            for i in range(0, n, block_size):
                end = min(i + block_size, n)
                block_slice = slice(i, end)
                
                if A_dense is not None:
                    # 使用稠密矩阵计算小块
                    residual = b[block_slice] - A_dense[block_slice, :].dot(x)
                    block_matrix = A_dense[block_slice, block_slice]
                    try:
                        x[block_slice] = np.linalg.solve(block_matrix, residual)
                    except np.linalg.LinAlgError:
                        # 如果矩阵奇异，使用最小二乘解
                        correction = np.linalg.lstsq(block_matrix, residual, rcond=None)[0]
                        x[block_slice] = correction
                else:
                    # 使用稀疏矩阵计算
                    residual = b[block_slice] - A[block_slice, :].dot(x)
                    block_matrix = A[block_slice, :][:, block_slice]
                    try:
                        correction = spsolve(block_matrix, residual)
                        x[block_slice] = correction
                    except Exception:
                        # 退回到雅可比迭代
                        diag = block_matrix.diagonal()
                        diag_inv = 1.0 / np.maximum(np.abs(diag), 1e-10) * np.sign(diag)
                        x[block_slice] += 0.5 * diag_inv * residual
        
        return x
    
    def _setup_grid_hierarchy(self, A, b):
        """设置多重网格层次结构（优化版）
        
        使用稀疏矩阵存储，减少内存使用
        
        Args:
            A: 稀疏系数矩阵
            b: 右端项向量
            
        Returns:
            tuple: (A_levels, b_levels, R_levels, P_levels)
        """
        print(f"创建多重网格层次结构，问题规模: {A.shape[0]}")
        
        # 计算适当的层数
        n = A.shape[0]
        max_levels = min(self.max_levels, int(np.log2(n/10)) + 1)
        max_levels = max(1, max_levels)
        
        # 初始化层次结构
        A_levels = [A]
        b_levels = [b]
        R_levels = []
        P_levels = []
        
        # 创建多层网格
        n_current = n
        for level in range(max_levels - 1):
            # 如果当前网格太小，不再继续粗化
            if n_current <= 100:
                break
                
            # 计算粗网格大小
            n_coarse = (n_current + 1) // 2
            
            try:
                # 创建限制和插值算子
                R = self._create_restriction_operator(n_current) 
                P = self._create_prolongation_operator(n_coarse, n_current)
                
                # 使用稀疏矩阵计算粗网格系数矩阵：A_coarse = R * A * P
                A_coarse = R.dot(A_levels[-1].dot(P))
                
                # 添加到层次结构
                A_levels.append(A_coarse)
                b_levels.append(np.zeros(n_coarse)) # 占位符
                R_levels.append(R)
                P_levels.append(P)
                
                # 更新当前网格大小
                n_current = n_coarse
            except Exception as e:
                print(f"创建网格层级{level+1}时出错: {e}")
                break
        
        print(f"多重网格层次结构创建完成，共{len(A_levels)}层")
        return A_levels, b_levels, R_levels, P_levels
    
    def _v_cycle(self, A_levels, b_levels, R_levels, P_levels, x):
        """执行V循环
        
        Args:
            A_levels: 各层次的系数矩阵列表
            b_levels: 各层次的右端项列表
            R_levels: 各层次的限制算子列表
            P_levels: 各层次的插值算子列表
            x: 初始猜测解
            
        Returns:
            numpy.ndarray: 改进后的解
        """
        n_levels = len(A_levels)
        x_levels = [None] * n_levels
        x_levels[0] = x.copy()
        
        # 下降阶段：从细网格到粗网格
        for level in range(n_levels - 1):
            # 前向光滑
            x_levels[level] = self._jacobi_smoothing(
                A_levels[level], x_levels[level], b_levels[level], self.pre_smoothing
            )
            
            # 计算残差
            residual = b_levels[level] - A_levels[level].dot(x_levels[level])
            
            # 将残差限制到粗网格
            b_levels[level+1] = R_levels[level].dot(residual)
            
            # 粗网格初始猜测为零
            x_levels[level+1] = np.zeros_like(b_levels[level+1])
        
        # 最粗网格求解
        try:
            if A_levels[-1].shape[0] <= 1000:
                # 对小问题使用直接求解
                x_levels[-1] = spsolve(A_levels[-1], b_levels[-1])
            else:
                # 对大问题使用迭代求解
                x_temp, success = cg(A_levels[-1], b_levels[-1], atol=1e-2, x0=np.zeros_like(b_levels[-1]), maxiter=200)
                if success == 0:
                    x_levels[-1] = x_temp
                else:
                    # 如果迭代法失败，使用光滑
                    x_levels[-1] = self._block_gauss_seidel(
                        A_levels[-1], x_levels[-1], b_levels[-1], 5, 10
                    )
        except Exception as e:
            print(f"粗网格求解错误: {e}")
            # 使用光滑作为备选方案
            x_levels[-1] = self._jacobi_smoothing(
                A_levels[-1], x_levels[-1], b_levels[-1], 20, 0.5
            )
        
        # 上升阶段：从粗网格到细网格
        for level in range(n_levels-2, -1, -1):
            # 插值误差校正
            correction = P_levels[level].dot(x_levels[level+1])
            x_levels[level] += correction
            
            # 后向光滑
            x_levels[level] = self._jacobi_smoothing(
                A_levels[level], x_levels[level], b_levels[level], self.post_smoothing
            )
        
        return x_levels[0]
    
    def solve(self, A, b):
        """求解线性方程组 Ax = b，针对大规模稀疏系统优化
        
        Args:
            A: 系数矩阵（稀疏矩阵）
            b: 右端项
            
        Returns:
            numpy.ndarray: 方程的解
        """
        n = A.shape[0]
        
        # 检查矩阵是否为稀疏矩阵
        if not isinstance(A, (csr_matrix, csc_matrix)):
            A = csr_matrix(A)
        
        # 对小型问题直接求解
        if n <= 500:
            try:
                return spsolve(A, b)
            except Exception as e:
                print(f"直接求解器失败: {e}，尝试迭代方法")
                # 继续使用多重网格
        
        # 设置多重网格层次结构
        A_levels, b_levels, R_levels, P_levels = self._setup_grid_hierarchy(A, b)
        
        # 初始猜测解
        x = np.zeros_like(b)
        
        # 多重网格迭代
        for iter_num in range(self.max_iterations):
            # V循环
            x_new = self._v_cycle(A_levels, b_levels, R_levels, P_levels, x)
            
            # 计算残差
            residual = b - A.dot(x_new)
            
            # 根据残差计算方式选择
            if self.residual_norm_type == 'relative':
                # 相对残差
                res_norm = np.linalg.norm(residual) / max(np.linalg.norm(b), 1e-10)
            else:
                # 绝对残差
                res_norm = np.linalg.norm(residual)
            
            # 检查收敛性（考虑最小迭代次数）
            if res_norm < self.tolerance and iter_num + 1 >= self.min_iterations:
                print(f"多重网格求解器在第{iter_num+1}次迭代收敛，残差: {res_norm:.6e}")
                return x_new
            
            # 更新解
            x = x_new
            
            # 每隔几次迭代输出一次残差
            if iter_num % 5 == 0 or iter_num + 1 == self.max_iterations:
                print(f"迭代 {iter_num+1}/{self.max_iterations}, 残差: {res_norm:.6e}")
        
        print(f"警告: 多重网格求解器未在{self.max_iterations}次迭代内收敛，最终残差: {res_norm:.6e}")
        
        # 返回最佳解
        return x
